import os

import numpy as np
from PIL import Image
from keras import Sequential
from keras.layers import Dense, Dropout
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer
import pandas as pd

# 加载路径
path = []
imagePaths = os.walk("train")
for (rootDir, dirNames, filenames) in imagePaths:
    for i in filenames:
        filename = "train/" + i
        path.append(filename)

#加载图片
data = []
labels = []
for imagePath in path:
    image = Image.open(imagePath)
    image = image.resize((28, 28))
    image = np.array(image).flatten()
    data.append(image)

    label = imagePath.split('/')[1].split('.')[0].split("_")[1]
    labels.append(label)

# 处理数据
data = np.array(data, dtype='float32') / 255.0
label = np.array(labels)

# 分割数据
(x_train, x_test, y_train, y_test) = train_test_split(data, labels, test_size=0.25, random_state=42)
lb = LabelBinarizer()
y_train = lb.fit_transform(y_train)
y_test = lb.fit_transform(y_test)

# 建立模型
model = Sequential()
model.add(Dense(512, input_shape=(x_train.shape[1],), activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(10, activation='softmax'))

# 编译模型
model.compile(loss="categorical_crossentropy", optimizer="rmsprop", metrics=["accuracy"])

# 训练模型
history = model.fit(x_train, y_train, batch_size=128, epochs=10, verbose=1, validation_split=0.1)








# =============================================================开始测试===============================================
# 加载路径
path = []
imagePaths = os.walk("test")
for (rootDir, dirNames, filenames) in imagePaths:
    for i in filenames:
        filename = "test/" + i
        path.append(filename)

#加载图片
test = []
test_l = []
ID = []
for imagePath in path:
    image = Image.open(imagePath)
    image = image.resize((28, 28))
    image = np.array(image).flatten()
    test.append(image)

    label = imagePath.split('/')[1].split('.')[0].split("_")[1]
    ID1 = imagePath.split('/')[1].split('.')[0].split("_")[0]
    ID.append(ID1)
    test_l.append(label)

# 处理数据
test = np.array(test, dtype='float32') / 255.0
test_l = np.array(test_l)

# 处理标签
lb = LabelBinarizer()
y_test = lb.fit_transform(test_l)

# 开始预测
y_pre=  model.predict(test)
y_pre = np.argmax(y_pre, axis=1)

pd.DataFrame(data={ID, y_pre}).to_csv('result.csv',index=False)